#include <iostream>
#include <windows.h>
#include <WtsApi32.h>
#include "config.h"
#include "log.h"
#include "filter_service.h"

#pragma comment(linker, "/subsystem:\"windows\" /entry:\"mainCRTStartup\"")

namespace {
const std::string PROCESS_NAME = "filter_agent";
const int MIN_ARG_LEN = 2;
const int CONFIG_ARG_INDEX = 1;
}  // namespace

int main(int argc, char **argv)
{
    HANDLE instance = CreateMutex(nullptr, FALSE, PROCESS_NAME.c_str());
    if (GetLastError() == ERROR_ALREADY_EXISTS) {
        CloseHandle(instance);
        std::cout << "process is running." << std::endl;
        return -1;
    }
    std::string configPath = "config.json";
    if (argc >= MIN_ARG_LEN) {
        configPath = std::string(argv[CONFIG_ARG_INDEX]);
    }
    FilterConfig config(configPath);

    Log::Init(config);
    auto service = std::make_shared<FilterService>();
    if (service == nullptr || !service->Init(config)) {
        LOG_ERR("init filter service failed.");
        CloseHandle(instance);
        return -1;
    }
    if (!service->Start()) {
        LOG_ERR("start filter service failed.");
        CloseHandle(instance);
        return -1;
    }
    HWND hWnd = CreateWindow("STATIC", "", WS_POPUP, 0, 0, 0, 0, nullptr, nullptr, GetModuleHandle(nullptr), nullptr);
    if (!hWnd || !WTSRegisterSessionNotification(hWnd, NOTIFY_FOR_THIS_SESSION)) {
        LOG_ERR("register failed. err: {}", GetLastError());
        CloseHandle(instance);
        return -1;
    }
    MSG msg{};
    while (GetMessage(&msg, hWnd, WM_WTSSESSION_CHANGE, WM_WTSSESSION_CHANGE)) {
        if (msg.message != WM_WTSSESSION_CHANGE) {
            continue;
        }
        switch (msg.wParam) {
            case WTS_SESSION_LOGON:
            case WTS_SESSION_UNLOCK:
            case WTS_REMOTE_CONNECT:
            case WTS_CONSOLE_CONNECT: {
                service->Start();
                break;
            }
            default: {
                service->Stop();
                break;
            }
        }
    }
    WTSUnRegisterSessionNotification(hWnd);
    DestroyWindow(hWnd);
    CloseHandle(instance);
    return 0;
}
